# -*- coding=utf-8 -*-
# 2022.0127
# @zlk
# primer design
import primer3
import os
import argparse

# DEFINING VARIABLES

parser = argparse.ArgumentParser()
parser.add_argument('-g', dest = 'g', required = True, 
                    help = 'path to genome'
                    )
parser.add_argument('-dhhp', dest = 'dhhp', required = True,
                    help = 'path to dhhp file'
                    )
parser.add_argument('-i', dest = 'i', required = True,
                    help = 'path to genome reigon, eg. chr11,0,10000'
                    )
parser.add_argument('-f', dest = 'f', default = './',
                    help = 'output folder'
                    )
parser.add_argument('-o' , '--output' , metavar = '' , dest = 'o' , default = 'output.primer.extracted' ,
                    help = 'path to output file in FASTA format, default = output.primer.extracted in Current running directory'
                    )
parser.add_argument('-n', '--number', metavar = '', dest = 'n', type = int, default = 10,
                    help = 'number of candidate primer pairs to pick, default = 10'
                    )
parser.add_argument('-k', '--flank', metavar = '', dest = 'k', type = int, default = 200,
                    help = 'flaning length of indel'
                    )
parser.add_argument('-s', '--short', metavar = '', dest = 's', type = int, default = 18,
                    help = 'shortest acceptable primer, default = 18'
                    )
parser.add_argument('-O', '--OPPORTUNE', metavar = '', dest = 'O', type = int, default = 20,
                    help = 'most acceptable primer, default = 20'
                    )
parser.add_argument('-l' , '--long' , metavar = '' , dest = 'l' , type = int , default = 24 ,
                    help = 'longest acceptable primer, default = 24'
                    )
parser.add_argument('-S' , '--SHORT' , metavar = '' , dest = 'S' , type = int , default = 70 ,
                    help = 'shortest product , default = 70'
                    )
parser.add_argument('-L' , '--LONG' , metavar = '' , dest = 'L' , type = int , default = 200 ,
                    help = 'longest product , default = 200'
                    )
parser.add_argument('-m' , '--mintemp' , metavar = '' , dest = 'm' , type = float , default = 50 ,
                    help = 'min Tm in celsius, default = 50'
                    )
parser.add_argument('-x' , '--maxtemp' , metavar = '' , dest = 'x' , type = float , default = 65 ,
                    help = 'max Tm in celsius, default = 65'
                    )
parser.add_argument('-M' , '--mingc' , metavar = '' , dest = 'M' , type = float , default = 35 ,
                    help = 'min GC percentage, default = 35'
                    )
parser.add_argument('-X' , '--maxgc' , metavar = '' , dest = 'X' , type = float , default = 65 ,
                    help = 'max GC percentage, default = 65'
                    )
parser.add_argument('-D' , '--tmdiff' , metavar = '' , dest = 'D' , type = float , default = 0.5 ,
                    help = 'accepted TM difference to form primer pair, default = 0.5'
                    )


args = parser.parse_args()
global_args = {
        'PRIMER_NUM_RETURN': args.n,
        'PRIMER_OPT_SIZE': args.O,
        'PRIMER_MIN_SIZE': args.s,
        'PRIMER_MAX_SIZE': args.l,
        'PRIMER_OPT_TM': 59.0,
        'PRIMER_MIN_TM': args.m,
        'PRIMER_MAX_TM': args.x,
        'PRIMER_MIN_GC': args.M,
        'PRIMER_MAX_GC': args.X,
        'PRIMER_THERMODYNAMIC_OLIGO_ALIGNMENT': 1,
        'PRIMER_MAX_POLY_X': 100,
        'PRIMER_INTERNAL_MAX_POLY_X': 100,
        'PRIMER_SALT_MONOVALENT': 50.0,
        'PRIMER_DNA_CONC': 50.0,
        'PRIMER_MAX_NS_ACCEPTED': 0,
        'PRIMER_MAX_SELF_ANY': 12,
        'PRIMER_MAX_SELF_END': 8,
        'PRIMER_PAIR_MAX_COMPL_ANY': 12,
        'PRIMER_PAIR_MAX_COMPL_END': 8,
        'PRIMER_PRODUCT_SIZE_RANGE': [args.S,args.L],
        'PRIMER_GC_CLAMP': 1,
        'RIMER_PAIR_MAX_DIFF_TM':args.D
}
# DEFINING MAIN PROGRAM



def pick_seq_of_indel(dhhp_file,genome_file,region,out_folder,flanking_length ):
    chr,start,end=region.split(',')[0],int(region.split(',')[1]),int(region.split(',')[2])
    
    #### pick indel sequence from genome########
    loop_genome=open(genome_file,'r')
    chr_seq,flag='',0
    for line in loop_genome:
        if line.startswith('>') :
            if line.strip()[1:]==chr:
                flag=1
            else:
                flag=0
        elif flag==1 :
            chr_seq+=line.strip()
    loop_genome.close()
    loop_dhhp=open(dhhp_file,'r')
    out_dhhp_dict={}
    out_indel_seq_file=open(out_folder+'/'+chr+'_'+str(start)+'_'+str(end)+'_indel.fasta','w')
    for line in loop_dhhp:
        line_list=line.strip().split('\t')
        if line_list[0]==chr and start<int(line_list[1])<end:
            out_dhhp_dict[line_list[0]+'_'+line_list[1]]=line_list[:8]+[line_list[-2]]
            if len(line_list[2])>4:
                out_seq=chr_seq[int(line_list[1]) - flanking_length : int(line_list[1])]+line_list[2]+chr_seq[int(line_list[1]) + 1 : int(line_list[1]) + flanking_length+1]
            elif len(line_list[3])>4:
                out_seq=chr_seq[int(line_list[1]) - flanking_length : int(line_list[1])]+line_list[3]+chr_seq[int(line_list[1]) + 1 : int(line_list[1]) + flanking_length+1]
            else:
                continue
            out_indel_seq_file.write('>'+chr+'_'+line_list[1]+'\n'+out_seq+'\n')
    loop_dhhp.close()
    out_indel_seq_file.close()
    ##### filter by blast ######
    make_db_command='makeblastdb -in '+ genome_file +' -dbtype nucl -parse_seqids -out '+ out_folder+'/db/'+genome_file
    blast_command= 'blastn -num_threads 30 -outfmt 6 -evalue 1e-5 -db '+ out_folder+'/db/'+genome_file+' -query '+out_folder+'/'+chr+'_'+str(start)+'_'+str(end)+'_indel.fasta -out ' + out_folder+'/indel_genome.blast'
    os.system(make_db_command)
    os.system(blast_command)
    indel_blast_file=open( out_folder+'/indel_genome.blast','r')
    count_dict={}
    for line in indel_blast_file:
        line_list=line.strip().split('\t')
        if int(line_list[3])<60:
            continue
        if line_list[0] not in count_dict.keys():
            count_dict[line_list[0]]=0
        count_dict[line_list[0]]+=1
    indel_blast_file.close()
    uniq_list=[]
    for key,value in count_dict.items():
        if value==1:
            uniq_list.append(key)
    out_indel_seq_file=open(out_folder+'/'+chr+'_'+str(start)+'_'+str(end)+'_indel.fasta','r')
    out_indel_seq_filter_file=open(out_folder+'/'+chr+'_'+str(start)+'_'+str(end)+'_indel_filter.fasta','w')
    out_dict={}
    for line in out_indel_seq_file:
        if line.startswith('>'):
            if line.strip()[1:] in uniq_list:
                flag=1
                out_indel_seq_filter_file.write(line)
                seq_id=line.strip()[1:]
                out_dict[seq_id]=''
            else:
                flag=0
        elif flag==1:
            out_dict[seq_id]+=line.strip()
            out_indel_seq_filter_file.write(line)
    out_indel_seq_file.close()
    return out_dict,out_dhhp_dict

def primer_main(in_dict,in_dhhp_dict,primer_file,global_args,flanking_length):
    out_file=open(primer_file,'w')
    out_file.write('Index\tChr\tPos\tRef\tAlt\tParent1\tParent2\tPool1\tPool2\tDhhp\tLeft_seq\tLeft_TM\tLeft_GC\tRight_seq\tRight_TM\tRight_GC\tProduct_size\n')
    index=0
    
    for key,value in in_dict.items():
        seq_args = {
        'SEQUENCE_ID': key,
        'SEQUENCE_TEMPLATE': value,
        'SEQUENCE_INCLUDED_REGION': [0,len(value)-1],
        'SEQUENCE_TARGET':[flanking_length,len(value)-2*flanking_length]
        }
        index+=1
        index1=0
        primer3_result = primer3.bindings.designPrimers(seq_args, global_args)
        for i in range(primer3_result['PRIMER_LEFT_NUM_RETURNED']):
            index1+=1
            out_file.write('primer_'+str(index)+'_'+str(index1)+'\t')
            for x in in_dhhp_dict[key]:
                out_file.write(x+'\t')
            out_file.write(primer3_result['PRIMER_LEFT_'+str(i)+'_SEQUENCE']+'\t')
            out_file.write(str(primer3_result['PRIMER_LEFT_'+str(i)+'_TM'])+'\t')
            out_file.write(str(primer3_result['PRIMER_LEFT_'+str(i)+'_GC_PERCENT'])+'\t')
            out_file.write(primer3_result['PRIMER_RIGHT_'+str(i)+'_SEQUENCE']+'\t')
            out_file.write(str(primer3_result['PRIMER_RIGHT_'+str(i)+'_TM'])+'\t')
            out_file.write(str(primer3_result['PRIMER_RIGHT_'+str(i)+'_GC_PERCENT'])+'\t')
            out_file.write(str(primer3_result['PRIMER_PAIR_'+str(i)+'_PRODUCT_SIZE'])+'\n')

if __name__=="__main__":
    indel_seq_dict,dhhp_dict=pick_seq_of_indel(args.dhhp,args.g,args.i,args.f ,args.k) 
    primer_main(indel_seq_dict,dhhp_dict,args.f+'/'+args.o,global_args,args.k)
 






